# Support Vector Machine (SVM) Classifier
from src.models.supmodel import SupModel
from sklearn.svm import SVC

class SVM(SupModel):
    """
    @brief 支持向量机分类器，支持线性和RBF核函数
    @details
    该类实现了支持向量机（SVM）分类器，支持线性核函数和RBF核函数。
    可通过max_iter参数设置最大迭代次数，默认为None表示不限制迭代次数。
    @exception ValueError: 如果指定的核函数不支持或max_iter不是整数
    @exception TypeError: 如果max_iter不是整数或None
    """
    def __init__(self, kernel='linear', max_iter=None):
        """
        @brief 初始化SVM分类器
        @param kernel 核函数类型，'linear'或'rbf'
        @param max_iter 最大迭代次数，默认不设置
        """
        if kernel not in ['linear', 'rbf']:
            raise ValueError("目前仅支持 'linear' 和 'rbf' 核函数")
        self.kernel = kernel
        self.max_iter = max_iter
        self.name = f"SVM_{self.kernel}_{self.max_iter if self.max_iter is not None else 'def'}"
        if max_iter is not None:
            self.model = SVC(kernel=self.kernel, max_iter=self.max_iter)
        else:
            self.model = SVC(kernel=self.kernel)

    def fit(self, X_train, y_train, sample_weight=None):
        """
        @brief 训练SVM模型
        @param X_train 训练特征
        @param y_train 训练标签
        @param sample_weight 样本权重（可选）
        """
        if sample_weight is not None:
            self.model.fit(X_train, y_train, sample_weight=sample_weight)
        else:
            self.model.fit(X_train, y_train)

    def predict(self, X_test):
        """
        @brief 对测试数据进行预测
        @param X_test 测试特征
        @return 预测标签
        """
        return self.model.predict(X_test)
